"""
准备闲聊语料
"""
import string
import chatbot.lib as lib
from tqdm import tqdm


def filter(pair):
    if pair[0][1].lower() in list(string.ascii_lowercase):
        return True
    elif pair[1][1].count("=") >= 2 and len(pair[1][0].split()) < 4:
        return True
    elif len(pair[0][0].strip()) == 0 or len(pair[1][0].strip()) == 0:
        return True


def prepar_xiaohuangji(by_word=False):
    path = "corpus/classify/origin_corpus/小黄鸡未分词.conv"
    input_path = "corpus/chatbot/input.txt"
    input_by_word_path = "corpus/chatbot/input_by_word.txt"
    target_path = "corpus/chatbot/target.txt"
    target_by_word_path = "corpus/chatbot/target_by_word.txt"
    f_input = open(input_by_word_path if by_word else input_path, "a", encoding='utf-8')
    f_target = open(target_by_word_path if by_word else target_path, "a", encoding='utf-8')

    qa_pair = []
    for line in tqdm(open(path, encoding='utf-8').readlines(), ascii=True, desc="处理小黄鸡语料："):
        if line.startswith("E"):
            continue
        else:
            line = line[1:].strip().lower()
            line_cuted = lib.cut(line, by_word=by_word)
            line_cuted = " ".join(line_cuted)
            if len(qa_pair) < 2:
                qa_pair.append([line_cuted, line])
            if len(qa_pair) == 2:
                assert len(qa_pair) == 2, "error"
                if filter(qa_pair):
                    qa_pair = []
                    continue
                f_input.write(qa_pair[0][0] + "\n")
                f_target.write(qa_pair[1][0] + "\n")
                qa_pair = []

    f_input.close()
    f_target.close()
